{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import yaml\n",
    "from torch.utils.data import DataLoader\n",
    "from dataset.PairKitti import PairKitti\n",
    "from dataset.PairCityscape import PairCityscape\n",
    "from models.balle2018.model import BMSHJ2018Model\n",
    "from models.balle2017.model import BLS2017Model\n",
    "from models.distributed_model import HyperPriorDistributedAutoEncoder, DistributedAutoEncoder\n",
    "from pytorch_msssim import ms_ssim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = \"configs/config.yaml\"\n",
    "with open(config_path, 'r') as stream:\n",
    "    config = yaml.load(stream, Loader=yaml.FullLoader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining Necessary Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Functions `get_bpp` and `get_distortion` are defined to get rate and distortion based on the model type. For both functions, the model output and the configuration are given as the input. Based on the type of the model which is obtained from `config`, `get_bpp` calculates rate, and `get_distortion` calculates distortion. `get_bpp` returns two values: the rates of the input image, side information, and the common information which are returned as one bpp value (for training and validation), and only the rate of the input image which is also returned as a bpp value (for testing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_bpp(model_out, config):  # Returns calculated bpp for train and test\n",
    "    alpha = config['alpha']\n",
    "    beta = config['beta']\n",
    "    if config['baseline_model'] == 'bmshj18':\n",
    "        if config['use_side_info']:  # If the side information (correlated image) has to be used\n",
    "            ''' \n",
    "            The loss function consists of:\n",
    "            Rate terms for input image (likelihoods), correlated image (y_likelihoods),\n",
    "            and the common information (w_likelihoods), hyperpriors for input image (z_likelihoods)\n",
    "            , hyperpriors for correlated image (z_likelihoods_cor).\n",
    "            Sum of these rate terms is returned as bpp, along with the actual bpp transmitted over the channel,\n",
    "            which consists only of likelihoods + z_likelihoods.\n",
    "            '''\n",
    "            x_recon, y_recon, likelihoods, y_likelihoods, z_likelihoods, z_likelihoods_cor, w_likelihoods = model_out\n",
    "            size_est = (-np.log(2) * x_recon.numel() / 3)\n",
    "            bpp = (torch.sum(torch.log(likelihoods)) + torch.sum(torch.log(z_likelihoods))) / size_est\n",
    "            transmitted_bpp = bpp.clone().detach()  # the real bpp value which is transmitted (for test)\n",
    "            bpp += alpha * (torch.sum(torch.log(y_likelihoods)) + torch.sum(torch.log(z_likelihoods_cor))) / size_est\n",
    "            bpp += beta * torch.sum(torch.log(w_likelihoods)) / size_est\n",
    "            return bpp, transmitted_bpp\n",
    "        else:  # The baseline implementation (Balle2018) without the side information\n",
    "            x_recon, likelihoods, z_likelihoods = model_out\n",
    "            size_est = (-np.log(2) * x_recon.numel() / 3)\n",
    "            bpp = (torch.sum(torch.log(likelihoods)) + torch.sum(torch.log(z_likelihoods))) / size_est\n",
    "            return bpp, bpp\n",
    "    elif config['baseline_model'] == 'bls17':\n",
    "        if config['use_side_info']:\n",
    "            x_recon, y_recon, likelihoods, y_likelihoods, w_likelihoods = model_out\n",
    "            size_est = (-np.log(2) * x_recon.numel() / 3)\n",
    "            bpp = torch.sum(torch.log(likelihoods)) / size_est\n",
    "            transmitted_bpp = bpp.clone().detach()  # the real bpp value which is transmitted (for test)\n",
    "            bpp += alpha * torch.sum(torch.log(y_likelihoods)) / size_est\n",
    "            bpp += beta * torch.sum(torch.log(w_likelihoods)) / size_est\n",
    "            return bpp, transmitted_bpp\n",
    "        else:\n",
    "            x_recon, likelihoods = model_out\n",
    "            size_est = (-np.log(2) * x_recon.numel() / 3)\n",
    "            bpp = torch.sum(torch.log(likelihoods)) / size_est\n",
    "            return bpp, bpp\n",
    "    return None\n",
    "\n",
    "\n",
    "def get_distortion(config, out, img, cor_img, mse):\n",
    "    distortion = None\n",
    "    alpha = config['alpha']\n",
    "    if config['use_side_info']:\n",
    "        ''' \n",
    "        The loss function consists of:\n",
    "        Distortion terms for input image (x_recon), and correlated image (x_cor_recon).\n",
    "        '''\n",
    "        x_recon, y_recon = out[0], out[1]\n",
    "        if config['distortion_loss'] == 'MS-SSIM':\n",
    "            distortion = (1 - ms_ssim(img.cpu(), x_recon.cpu(), data_range=1.0, size_average=True,\n",
    "                                      win_size=7))\n",
    "            distortion += alpha * (1 - ms_ssim(cor_img.cpu(), y_recon.cpu(), data_range=1.0, size_average=True,\n",
    "                                               win_size=7))\n",
    "        elif config['distortion_loss'] == 'MSE':\n",
    "            distortion = mse(img, x_recon)\n",
    "            distortion += alpha * mse(cor_img, y_recon)\n",
    "    else:\n",
    "        x_recon = out[0]\n",
    "        if config['distortion_loss'] == 'MS-SSIM':\n",
    "            distortion = (1 - ms_ssim(img.cpu(), x_recon.cpu(), data_range=1.0, size_average=True,\n",
    "                                      win_size=7))\n",
    "        elif config['distortion_loss'] == 'MSE':\n",
    "            distortion = mse(img, x_recon)\n",
    "\n",
    "    return distortion\n",
    "\n",
    "\n",
    "def save_image(x_recon, x, path, name):\n",
    "    img_recon = np.clip((x_recon * 255).squeeze().cpu().numpy(), 0, 255)\n",
    "    img = np.clip((x * 255).squeeze().cpu().numpy(), 0, 255)\n",
    "    img_recon = np.transpose(img_recon, (1, 2, 0)).astype('uint8')\n",
    "    img_final = Image.fromarray(np.concatenate((img, img_recon), axis=1), 'RGB')\n",
    "    if not os.path.exists(path):\n",
    "        os.makedirs(path)\n",
    "    img_final.save(os.path.join(path, name + '.png'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the Datasets and the Dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset initialization\n",
    "path = config['dataset_path']\n",
    "resize = tuple(config['resize'])\n",
    "if config['dataset_name'] == 'KITTI':\n",
    "    train_dataset = PairKitti(path=path, set_type='train', resize=resize)\n",
    "    val_dataset = PairKitti(path=path, set_type='val', resize=resize)\n",
    "    test_dataset = PairKitti(path=path, set_type='test', resize=resize)\n",
    "elif config['dataset_name'] == 'Cityscape':\n",
    "    train_dataset = PairCityscape(path=path, set_type='train', resize=resize)\n",
    "    val_dataset = PairCityscape(path=path, set_type='val', resize=resize)\n",
    "    test_dataset = PairCityscape(path=path, set_type='test', resize=resize)\n",
    "else:\n",
    "    raise Exception(\"Dataset not found\")\n",
    "\n",
    "batch_size = config['train_batch_size']\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=3)\n",
    "val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=True, num_workers=3)\n",
    "test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False, num_workers=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model initialization\n",
    "'''\n",
    "We provide the option of two baseline models.\n",
    "1) The Balle2017 model (bls17).\n",
    "2) The Balle2018 model (bmshj18), which uses scale hyperpriors.\n",
    "'''\n",
    "with_side_info = config['use_side_info']\n",
    "model_class = None\n",
    "if config['baseline_model'] == 'bmshj18':\n",
    "    if with_side_info:\n",
    "        model_class = HyperPriorDistributedAutoEncoder\n",
    "    else:\n",
    "        model_class = BMSHJ2018Model\n",
    "elif config['baseline_model'] == 'bls17':\n",
    "    if with_side_info:\n",
    "        model_class = DistributedAutoEncoder\n",
    "    else:\n",
    "        model_class = BLS2017Model\n",
    "\n",
    "model = model_class(num_filters=config['num_filters'])\n",
    "model = model.cuda() if config['cuda'] else model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training and Testing Preparations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], amsgrad=True)\n",
    "if config['load_weight']:\n",
    "    checkpoint = torch.load(config['weight_path'])\n",
    "    model.load_state_dict(checkpoint['model_state_dict'])\n",
    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, min_lr=1e-7)\n",
    "experiment_name = model_class.__name__ + '_' + str(train_dataset) + '_' + config['distortion_loss'] + '_lambda:' + \\\n",
    "                      str(config['lambda'])\n",
    "\n",
    "print('Experiment: ', experiment_name)\n",
    "\n",
    "weight_folder = None\n",
    "if config['save_weights']:\n",
    "    weight_folder = os.path.join(config['save_output_path'], 'weight')\n",
    "    if not os.path.exists(weight_folder):\n",
    "        os.makedirs(weight_folder)\n",
    "\n",
    "mse = torch.nn.MSELoss(reduction='mean')\n",
    "mse = mse.cuda() if config['cuda'] else mse\n",
    "lmbda = config['lambda']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(config['epochs']):\n",
    "    min_val_loss = None\n",
    "    model.train()\n",
    "    for i, data in enumerate(iter(train_loader)):\n",
    "        img, cor_img, _, _ = data\n",
    "        img = img.cuda().float() if config['cuda'] else img.float()\n",
    "        cor_img = cor_img.cuda().float() if config['cuda'] else cor_img.float()\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        if with_side_info:\n",
    "            out = model(img, cor_img)\n",
    "        else:\n",
    "            out = model(img)\n",
    "        bpp, _ = get_bpp(out, config)\n",
    "\n",
    "        distortion = get_distortion(config, out, img, cor_img, mse)\n",
    "\n",
    "        loss = lmbda * distortion * (255 ** 2) + bpp  # multiplied by (255 ** 2) for distortion scaling\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # Validation\n",
    "    model.eval()\n",
    "    val_loss = []\n",
    "    val_mse = []\n",
    "    val_msssim = []\n",
    "    val_bpp = []\n",
    "    val_transmitted_bpp = []\n",
    "    val_distortion = []\n",
    "    with torch.no_grad():\n",
    "        for i, data in enumerate(iter(val_loader)):\n",
    "            # img = input image, cor_img = side information/correlated image (designated y in the paper)\n",
    "            img, cor_img, _, _ = data\n",
    "            img = img.cuda().float() if config['cuda'] else img.float()\n",
    "            cor_img = cor_img.cuda().float() if config['cuda'] else cor_img.float()\n",
    "\n",
    "            if with_side_info:\n",
    "                out = model(img, cor_img)\n",
    "            else:\n",
    "                out = model(img)\n",
    "            bpp, transmitted_bpp = get_bpp(out, config)\n",
    "\n",
    "            x_recon = out[0]\n",
    "            mse_dist = mse(img, x_recon)\n",
    "            msssim = 1 - ms_ssim(img.clone().cpu(), x_recon.clone().cpu(), data_range=1.0, size_average=True,\n",
    "                                 win_size=7)\n",
    "            msssim_db = -10 * np.log10(msssim)\n",
    "\n",
    "            distortion = get_distortion(config, out, img, cor_img, mse)\n",
    "\n",
    "            loss = lmbda * distortion * (255 ** 2) + bpp  # multiplied by (255 ** 2) for distortion scaling\n",
    "\n",
    "            val_mse.append(mse_dist.item())\n",
    "            val_bpp.append(bpp.item())\n",
    "            val_transmitted_bpp.append(transmitted_bpp.item())\n",
    "            val_loss.append(loss.item())\n",
    "            val_msssim.append(msssim_db.item())\n",
    "            val_distortion.append(distortion.item())\n",
    "\n",
    "    val_loss_to_track = sum(val_loss) / len(val_loss)\n",
    "    scheduler.step(val_loss_to_track)\n",
    "\n",
    "    # Verbose\n",
    "    if config['verbose_period'] > 0 and (epoch + 1) % config['verbose_period'] == 0:\n",
    "        tracking = ['Epoch {}:'.format(epoch + 1),\n",
    "                    'Loss = {:.4f},'.format(val_loss_to_track),\n",
    "                    'BPP = {:.4f},'.format(sum(val_bpp) / len(val_bpp)),\n",
    "                    'Distortion = {:.4f},'.format(sum(val_distortion) / len(val_distortion)),\n",
    "                    'Transmitted BPP = {:.4f},'.format(sum(val_transmitted_bpp) / len(val_transmitted_bpp)),\n",
    "                    'PSNR = {:.4f},'.format(10 * np.log10(255 ** 2 / (sum(val_mse) / (len(val_mse))))),\n",
    "                    'MS-SSIM = {:.4f}'.format(sum(val_msssim) / len(val_msssim))]\n",
    "        print(\" \".join(tracking))\n",
    "        \n",
    "    # Save weights\n",
    "    if config['save_weights']:\n",
    "        if min_val_loss is None or min_val_loss > val_loss_to_track:\n",
    "            min_val_loss = val_loss_to_track\n",
    "            save_path = os.path.join(weight_folder, experiment_name + '.pt')\n",
    "            torch.save({\n",
    "                \"model_state_dict\": model.state_dict(),\n",
    "                \"optimizer_state_dict_{:}\": optimizer.state_dict(),\n",
    "            }, save_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_path = os.path.join(config['save_output_path'], 'results')\n",
    "if not os.path.exists(results_path):\n",
    "    os.makedirs(results_path)\n",
    "names = [\"Image Number\", \"Loss\", \"BPP\", \"PSNR\", \"MS-SSIM\"]\n",
    "cols = dict()\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for i, data in enumerate(iter(test_loader)):\n",
    "        img, cor_img, _, _ = data\n",
    "        img = img.cuda().float() if config['cuda'] else img.float()\n",
    "        cor_img = cor_img.cuda().float() if config['cuda'] else cor_img.float()\n",
    "\n",
    "        if with_side_info:\n",
    "            out = model(img, cor_img)\n",
    "        else:\n",
    "            out = model(img)\n",
    "        bpp, transmitted_bpp = get_bpp(out, config)\n",
    "\n",
    "        x_recon = out[0]\n",
    "        mse_dist = mse(img, x_recon)\n",
    "        msssim = 1 - ms_ssim(img.clone().cpu(), x_recon.clone().cpu(), data_range=1.0, size_average=True,\n",
    "                             win_size=7)\n",
    "        msssim_db = -10 * np.log10(msssim)\n",
    "\n",
    "        distortion = get_distortion(config, out, img, cor_img, mse)\n",
    "        loss = lmbda * distortion * (255 ** 2) + bpp  # multiplied by (255 ** 2) for distortion scaling\n",
    "\n",
    "        vals = [str(i)] + ['{:.8f}'.format(x) for x in [loss.item(),\n",
    "                                                        transmitted_bpp.item(),\n",
    "                                                        10 * np.log10(255 ** 2 / mse_dist.item()),\n",
    "                                                        msssim_db.item()]]\n",
    "\n",
    "        for (name, val) in zip(names, vals):\n",
    "            if name not in cols:\n",
    "                cols[name] = []\n",
    "            cols[name].append(val)\n",
    "\n",
    "        if config['save_image']:\n",
    "            save_image(x_recon[0], img, os.path.join(results_path, 'test_images'), str(i))\n",
    "\n",
    "    df = pd.DataFrame.from_dict(cols)\n",
    "    df.to_csv(os.path.join(results_path, experiment_name + '.csv'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}